import os
import time
import uuid
import json
import random
from autogen import AssistantAgent
from data_loader import load_topic_reports
from progress_manager import (
    extract_stance,
    save_progress_with_results,
    load_progress_with_results,
    backup_results,
    cleanup_progress,
    should_save_progress,
    parse_progress_arguments,
    get_progress_config
)
import re

def extract_stance_category(text, split_by="【立场】"):
    """从文本中提取立场类别（如：一般支持、完全反对等）"""
    if split_by in text:
        stance_text = text.split(split_by, 1)[1].strip()
        
        # 使用正则表达式找到下一个【】括起来的部分
        next_split = re.search(r'【.*?】', stance_text)
        if next_split:
            stance_text = stance_text[:next_split.start()].strip()
        
        # 定义可能的立场类别
        stance_categories = [
            "完全反对", "一般反对", "中立", "一般支持", "完全支持"
        ]
        
        # 在文本中查找立场类别
        for category in stance_categories:
            if category in stance_text:
                # 找到立场类别后，提取该类别（可能包含冒号等）
                # 使用正则表达式匹配立场类别及其可能的后续内容
                pattern = rf"({category})[：:]*\s*(.*?)(?=\n|$|。|，|；)"
                match = re.search(pattern, stance_text)
                if match:
                    return match.group(1).strip()  # 只返回立场类别
        
        # 如果没有找到明确的立场类别，尝试从文本中推断
        stance_text_lower = stance_text.lower()
        if any(word in stance_text_lower for word in ["完全反对", "坚决反对", "强烈反对"]):
            return "完全反对"
        elif any(word in stance_text_lower for word in ["一般反对", "不支持"]):
            return "一般反对"
        elif any(word in stance_text_lower for word in ["中立", "中立", "不偏不倚"]):
            return "中立"
        elif any(word in stance_text_lower for word in ["一般支持", "倾向于支持"]):
            return "一般支持"
        elif any(word in stance_text_lower for word in ["完全支持", "坚决支持", "强烈支持"]):
            return "完全支持"
        
        # 如果还是没找到，返回原始文本的前20个字符作为参考
        return stance_text[:20].strip()
    
    return "立场未找到"

# 设置随机种子
random.seed(42)

# 大五人格量表定义 - 基于BFI-2量表
BIG5_ITEMS = {
    "Extraversion": {
        "high": [
            "I am outgoing, sociable.",  # 1
            "I have an assertive personality.",  # 6
            "I am dominant, acts as a leader.",  # 21
            "I am full of energy.",  # 41
            "I am talkative.",  # 46
            "I show a lot of enthusiasm."  # 56
        ],
        "low": [
            "I rarely feel excited or eager.",  # 11R
            "I tend to be quiet.",  # 16R
            "I am sometimes shy, introverted.",  # 31R
            "I find it hard to influence people.",  # 36R
            "I am less active than other people.",  # 26R
            "I prefer to have others take charge."  # 51R
        ]
    },
    "Agreeableness": {
        "high": [
            "I am compassionate, has a soft heart.",  # 2
            "I am respectful, treats others with respect.",  # 7
            "I have a forgiving nature.",  # 27
            "I am helpful and unselfish with others.",  # 32
            "I am polite, courteous to others.",  # 52
            "I assume the best about people."  # 57
        ],
        "low": [
            "I tend to find fault with others.",  # 12R
            "I feel little sympathy for others.",  # 17R
            "I start arguments with others.",  # 22R
            "I am sometimes rude to others.",  # 37R
            "I am suspicious of others' intentions.",  # 42R
            "I can be cold and uncaring."  # 47R
        ]
    },
    "Conscientiousness": {
        "high": [
            "I am dependable, steady.",  # 13
            "I am systematic, likes to keep things in order.",  # 18
            "I keep things neat and tidy.",  # 33
            "I am efficient, gets things done.",  # 38
            "I am reliable, can always be counted on.",  # 43
            "I am persistent, works until the task is finished."  # 53
        ],
        "low": [
            "I tend to be disorganized.",  # 3R
            "I tend to be lazy.",  # 8R
            "I have difficulty getting started on tasks.",  # 23R
            "I can be somewhat careless.",  # 28R
            "I leave a mess, doesn't clean up.",  # 48R
            "I sometimes behave irresponsibly."  # 58R
        ]
    },
    "Neuroticism": {
        "high": [
            "I am moody, has up and down mood swings.",  # 14
            "I can be tense.",  # 19
            "I worry a lot.",  # 34
            "I often feel sad.",  # 39
            "I tend to feel depressed, blue.",  # 54
            "I am temperamental, gets emotional easily."  # 59
        ],
        "low": [
            "I am relaxed, handles stress well.",  # 4R
            "I stay optimistic after experiencing a setback.",  # 9R
            "I feel secure, comfortable with self.",  # 24R
            "I am emotionally stable, not easily upset.",  # 29R
            "I keep their emotions under control.",  # 44R
            "I rarely feel anxious or afraid."  # 49R
        ]
    },
    "Openness": {
        "high": [
            "I am curious about many different things.",  # 10
            "I am inventive, finds clever ways to do things.",  # 15
            "I am fascinated by art, music, or literature.",  # 20
            "I am complex, a deep thinker.",  # 40
            "I value art and beauty.",  # 35
            "I am original, comes up with new ideas."  # 60
        ],
        "low": [
            "I have few artistic interests.",  # 5R
            "I avoid intellectual, philosophical discussions.",  # 25R
            "I have little creativity.",  # 30R
            "I have difficulty imagining things.",  # 45R
            "I think poetry and plays are boring.",  # 50R
            "I have little interest in abstract ideas."  # 55R
        ]
    }
}

# 实验配置
EXPERIMENT_CONFIG = {
    "topics_per_experiment": 12,  # 每个实验使用的话题数量（3个手动选择 + 3个随机选择）
    "llm_config": {
        "config_list": [{
            # "model": "Qwen3-32B",
            # "base_url": "http://172.17.65.44:8000/v1",
            "model": "qwen3-4b",
            "base_url": "http://172.17.65.41:10125/v1/",
            "api_key": "",
        }],
        "temperature": 0.7,
    },
    "progress": {
        "enabled": True,
        "save_interval": 5,
        "progress_file": "big5_prior_progress.json",
        "auto_cleanup": True,
        "resume_from_progress": True,
    }
}

def generate_big5_profile(personality_traits):
    """
    根据大五人格特征生成profile描述
    
    Args:
        personality_traits: dict, 包含五个维度的特征，如 {"Extraversion": "high", "Agreeableness": "low", ...}
    
    Returns:
        str: 生成的profile描述
    """
    profile_parts = ["I am a normal person described by the Big Five Personality Traits."]
    
    for trait, level in personality_traits.items():
        if trait in BIG5_ITEMS and level in ["high", "low"]:
            descriptions = BIG5_ITEMS[trait][level]
            profile_parts.append(f"In terms of {trait}, I show {level} levels: {', '.join(descriptions)}")
    
    return ". ".join(profile_parts) + "."

def create_prior_experiment_combinations():
    """
    创建先验实验组合：大五人格特征 × 话题
    """
    combinations = []
    
    # 加载话题数据
    topic_path = "/data01/public/chenmiaokun/everything/cognitive_manipulation/datasets/opinion_pro/procons.pkl"
    topic_reports = load_topic_reports(topic_path, max_pro=3, max_con=3)
    
    # 选择话题
    available_topics = list(topic_reports.keys())
    
    # 手动选择三个话题
    manual_topics = [
        "Should Adults Have the Right to Carry a Concealed Handgun?",
        "Should the United States Return to a Gold Standard?",
        "Should Fighting Be Allowed in Hockey?"
    ]
    
    # 验证手动选择的话题是否在可用话题列表中
    valid_manual_topics = []
    for topic in manual_topics:
        if topic in available_topics:
            valid_manual_topics.append(topic)
        else:
            print(f"⚠️ 警告：话题 '{topic}' 不在可用话题列表中")
            print(f"可用话题：{available_topics[:5]}...")  # 显示前5个话题作为参考
    
    # 从剩余话题中随机选择额外的话题
    remaining_topics = [topic for topic in available_topics if topic not in valid_manual_topics]
    additional_topics_count = EXPERIMENT_CONFIG["topics_per_experiment"] - len(valid_manual_topics)
    
    if additional_topics_count > 0 and remaining_topics:
        additional_topics = random.sample(remaining_topics, min(additional_topics_count, len(remaining_topics)))
        selected_topics = valid_manual_topics + additional_topics
    else:
        selected_topics = valid_manual_topics
    
    print(f"手动选择的话题：{valid_manual_topics}")
    print(f"随机选择的话题：{selected_topics[len(valid_manual_topics):]}")
    print(f"最终选择的话题：{selected_topics}")
    
    # 生成所有可能的大五人格组合
    personality_dimensions = ["Extraversion", "Agreeableness", "Conscientiousness", "Neuroticism", "Openness"]
    personality_levels = ["high", "low"]
    
    # 为每个维度生成高/低两种情况的profile
    for ext in personality_levels:
        for agr in personality_levels:
            for con in personality_levels:
                for neu in personality_levels:
                    for ope in personality_levels:
                        personality_traits = {
                            "Extraversion": ext,
                            "Agreeableness": agr,
                            "Conscientiousness": con,
                            "Neuroticism": neu,
                            "Openness": ope
                        }
                        
                        profile = generate_big5_profile(personality_traits)
                        
                        # 为每个话题创建组合
                        for topic in selected_topics:
                            combinations.append({
                                "profile": profile,
                                "personality_traits": personality_traits,
                                "topic": topic,
                                "topic_report": topic_reports[topic]
                            })
    
    return combinations

def run_single_prior_experiment(combination, experiment_index, total_experiments):
    """
    运行单个先验实验
    
    Args:
        combination: dict, 实验组合
        experiment_index: int, 实验索引
        total_experiments: int, 总实验数
    
    Returns:
        dict: 实验结果
    """
    profile = combination["profile"]
    personality_traits = combination["personality_traits"]
    topic = combination["topic"]
    topic_report = combination["topic_report"]
    
    print(f"\n>>> 运行先验实验 #{experiment_index+1}/{total_experiments}")
    print(f"   profile: {profile}")
    print(f"   人格特征: {personality_traits}")
    print(f"   话题: {topic}")
    
    try:
        # 生成初始立场
        print("   生成初始立场...")
        
        normal_reader = AssistantAgent(
            name="NormalReader",
            system_message=(
                f"你是一个普通人，具有以下人格特征：{profile}\n\n"
                "请思考以下话题，并根据你的人格特征形成观点。\n\n"
                "请按照以下格式回答：\n"
                "【关键信息提取】：从话题中提取关键信息\n"
                "【个人分析】：根据你的人格特征进行分析\n"
                "【立场】：选择并说明你的立场：\n"
                "- 完全反对\n"
                "- 一般反对\n"
                "- 中立\n"
                "- 一般支持\n"
                "- 完全支持\n"
                "/no_think"
            ),
            llm_config=EXPERIMENT_CONFIG["llm_config"]
        )
        
        initial_response = normal_reader.generate_reply(
            messages=[{"role": "user", "content": f"话题：{topic}\n相关资料：{topic_report}"}]
        )
        
        if initial_response is None:
            initial_response = "No response"
        elif isinstance(initial_response, dict):
            initial_response = initial_response.get('content', "No response")
        else:
            initial_response = str(initial_response)
        
        initial_stance = extract_stance_category(initial_response, '【立场】')
        print(f"   初始立场: {initial_stance}")
        
        return {
            "id": str(uuid.uuid4()),
            "timestamp": time.strftime("%Y-%m-%d %H:%M:%S"),
            "personality_traits": personality_traits,
            "profile": profile,
            "topic": topic,
            "initial_stance": initial_response,
            "extracted_stance": initial_stance,
        }
        
    except Exception as e:
        print(f"   ❌ 实验执行失败: {e}")
        return {
            "id": str(uuid.uuid4()),
            "timestamp": time.strftime("%Y-%m-%d %H:%M:%S"),
            "personality_traits": personality_traits,
            "profile": profile,
            "topic": topic,
            "initial_stance": "Execution error",
            "extracted_stance": "Execution error",
        }

def analyze_prior_results(results):
    """
    分析先验实验结果
    
    Args:
        results: list, 实验结果列表
    
    Returns:
        dict: 分析结果
    """
    analysis = {
        "total_experiments": len(results),
        "topic_analysis": {},
        "personality_analysis": {},
        "stance_distribution": {
            "完全反对": 0,
            "一般反对": 0,
            "中立": 0,
            "一般支持": 0,
            "完全支持": 0
        }
    }
    
    # 按话题分析
    topics = set(result["topic"] for result in results)
    for topic in topics:
        topic_results = [r for r in results if r["topic"] == topic]
        analysis["topic_analysis"][topic] = {
            "total": len(topic_results),
            "stance_distribution": {
                "完全反对": 0,
                "一般反对": 0,
                "中立": 0,
                "一般支持": 0,
                "完全支持": 0
            },
            "personality_stance_correlation": {}
        }
        
        for result in topic_results:
            stance = result["extracted_stance"]
            if stance in analysis["topic_analysis"][topic]["stance_distribution"]:
                analysis["topic_analysis"][topic]["stance_distribution"][stance] += 1
                analysis["stance_distribution"][stance] += 1
            
            # 分析人格特征与立场的关系
            for trait, level in result["personality_traits"].items():
                key = f"{trait}_{level}"
                if key not in analysis["topic_analysis"][topic]["personality_stance_correlation"]:
                    analysis["topic_analysis"][topic]["personality_stance_correlation"][key] = {
                        "total": 0,
                        "stance_distribution": {
                            "完全反对": 0,
                            "一般反对": 0,
                            "中立": 0,
                            "一般支持": 0,
                            "完全支持": 0
                        }
                    }
                
                analysis["topic_analysis"][topic]["personality_stance_correlation"][key]["total"] += 1
                if stance in analysis["topic_analysis"][topic]["personality_stance_correlation"][key]["stance_distribution"]:
                    analysis["topic_analysis"][topic]["personality_stance_correlation"][key]["stance_distribution"][stance] += 1
    
    # 按人格特征分析
    personality_dimensions = ["Extraversion", "Agreeableness", "Conscientiousness", "Neuroticism", "Openness"]
    for trait in personality_dimensions:
        analysis["personality_analysis"][trait] = {
            "high": {"total": 0, "stance_distribution": {"完全反对": 0, "一般反对": 0, "中立": 0, "一般支持": 0, "完全支持": 0}},
            "low": {"total": 0, "stance_distribution": {"完全反对": 0, "一般反对": 0, "中立": 0, "一般支持": 0, "完全支持": 0}}
        }
        
        for result in results:
            level = result["personality_traits"][trait]
            stance = result["extracted_stance"]
            
            analysis["personality_analysis"][trait][level]["total"] += 1
            if stance in analysis["personality_analysis"][trait][level]["stance_distribution"]:
                analysis["personality_analysis"][trait][level]["stance_distribution"][stance] += 1
    
    return analysis

def main():
    """主函数"""
    # 解析命令行参数
    args = parse_progress_arguments()
    
    # 获取进度配置
    progress_config = get_progress_config(args)
    EXPERIMENT_CONFIG["progress"] = progress_config
    
    print("🔬 大五人格先验实验开始")
    print(f"实验配置：{EXPERIMENT_CONFIG}")
    
    # 创建实验组合
    combinations = create_prior_experiment_combinations()
    print(f"总实验数：{len(combinations)}")
    
    # 进度恢复处理
    progress_file = progress_config.get("progress_file", "big5_prior_progress.json")
    start_index, existing_results = load_progress_with_results(progress_file)
    results = existing_results  # 从已保存的结果开始
    
    print(f"🔄 从进度恢复: 已完成 {len(existing_results)} 个实验，从第 {start_index} 个开始")
    
    # 运行实验
    for i, combination in enumerate(combinations[start_index:], start=start_index):
        result = run_single_prior_experiment(combination, i, len(combinations))
        results.append(result)
        
        # 进度保存（包含已完成的实验结果）
        if should_save_progress(i, len(combinations), progress_config.get("save_interval", 5)):
            save_progress_with_results(i + 1, len(combinations), results, progress_file)
            
            # 如果启用了结果备份
            if progress_config.get("backup_results", False):
                backup_results(results, f"prior_backup_{int(time.time())}.json")
        
        # 限制实验数量（用于测试）
        # if i > 1:
        #     break
    
    # 分析结果
    analysis = analyze_prior_results(results)
    
    # 保存结果
    os.makedirs("results", exist_ok=True)
    json_path = f"results/big5_prior_experiment_{EXPERIMENT_CONFIG['llm_config']['config_list'][0]['model']}_{int(time.time())}.json"
    
    output_data = {
        "experiments": results,
        "analysis": analysis
    }
    
    with open(json_path, "w", encoding="utf-8") as f:
        json.dump(output_data, f, ensure_ascii=False, indent=2)
    
    print(f"\n✅ 先验实验完成，结果已保存：{json_path}")
    
    # 清理进度文件
    if progress_config.get("enabled", True) and progress_config.get("auto_cleanup", True):
        progress_file = progress_config.get("progress_file", "big5_prior_progress.json")
        cleanup_progress(progress_file)
    
    # 输出分析摘要
    print(f"\n📊 先验实验结果摘要：")
    print(f"   总实验数: {analysis['total_experiments']}")
    
    print(f"\n📈 总体立场分布：")
    for stance, count in analysis["stance_distribution"].items():
        percentage = (count / analysis["total_experiments"] * 100) if analysis["total_experiments"] > 0 else 0
        print(f"   {stance}: {count} ({percentage:.1f}%)")
    
    print(f"\n🎯 按话题分析：")
    for topic, topic_analysis in analysis["topic_analysis"].items():
        print(f"   话题：{topic}")
        print(f"     总实验数: {topic_analysis['total']}")
        for stance, count in topic_analysis["stance_distribution"].items():
            if count > 0:
                percentage = (count / topic_analysis["total"] * 100)
                print(f"     {stance}: {count} ({percentage:.1f}%)")
    
    print(f"\n🧠 按人格特征分析：")
    for trait, trait_analysis in analysis["personality_analysis"].items():
        print(f"   {trait}:")
        for level, level_analysis in trait_analysis.items():
            if level_analysis["total"] > 0:
                print(f"     {level}: {level_analysis['total']} 个实验")
                for stance, count in level_analysis["stance_distribution"].items():
                    if count > 0:
                        percentage = (count / level_analysis["total"] * 100)
                        print(f"       {stance}: {count} ({percentage:.1f}%)")

if __name__ == "__main__":
    main() 